clc;clear all;close all;
%main file to solve for the competitive equilibrium before the reform
global alpha betta depr sigmaC sigmaL gassetrat tauL tauK taumax lam
global phi popweight N govexp govery
global e_real omega kg

alpha = 0.394;   %capital share
betta = 0.96;   %discount factor
depr = 0.08;    %depreciation rate
sigmaC = 2;   %coefficient of relative risk aversion
sigmaL = 3;   %inverse Frisch elasticity

% tauL = 0.23;  %tax rates under the status quo policy 
% tauK = 0.57;
tauL=0.214 %Average of U.S. labour income tax rates taken from Trabandt and Uhlig 2012, period 2001-2010.
tauK=0.401  %Average of U.S. capital income tax with depreciation allowance taken from Trabandt and Uhlig 2012, period 2001-2010
taumax = tauK;  %max. capital tax rate
gassetrat = -0.668;  %initial government debt as a share of output

e_real=0.245

%Find omega (weight of leisure) such that lss=0.245439 at the steady state in the homogeneous agent case
N = 1;
phi =1; %labor productivity is normalized to 1
[gg,info] = fzero('find_omega_N1',0.12)%,0,1) %finds g
kss=(((1/betta-1)/(1-tauK)+depr)/alpha/e_real^(1-alpha))^(1/(alpha-1))
css=kss^alpha*e_real^(1-alpha)-gg-depr*kss
omegaN1=(1-alpha)*(kss/e_real)^alpha*(1-tauL)*css^(-sigmaC)/e_real^(sigmaL)
gg/kss^alpha/e_real^(1-alpha) %check g/y

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%HETEROGENOUS ECONOMY BEFORE THE REFORM

N=2

popweight = [.5 .5];  %population shares
phi = [1.1 1]'; %labor productivity
phi = phi/(popweight*phi); %normalize average labor productivity to 1
N = length(popweight);
lam=0.54 %c_j/c_1

options = optimset('MaxFunEvals',1000000,'MaxIter',1000,'TolFun',1e-10,'TolX',1e-10,'Disp','iter');
x0=[0.1 -1 omegaN1 -0.2 0.1] %g kj omega kg depr
lb=[0.01 -3 100 -1 0.02]
ub=[0.3 1 1000 -0.01 0.2]
govery=0.2
% x0=[0.1 -0.5 omegaN1 -0.2] %g kj omega kg
% lb=[0.01 -3 200 -1]
% ub=[0.3 0 300 0]
[GE,info] = lsqnonlin('find_gk',x0,lb,ub,options);
% opts1=[1e-7 1 1 1e-6]
% [GE,info] = broydn('find_glam',x0,opts1);
govexp = GE(1)
kj2N = [GE(2:N)]
omega=GE(N+1)
e_real=0.245
kg=GE(N+2)
depr=GE(N+3)

kss=(((1/betta-1)/(1-tauK)+depr)/alpha/e_real^(1-alpha))^(1/(alpha-1)); %the Euler at ss gives aggregate capital
css=(kss^alpha*e_real^(1-alpha)-govexp-depr*kss)/(popweight(1)+popweight(2:N)*lam'); %the resource constraint gives c1

%find l1 from optimality conditions of ratio of hours of different groups
%and using that average hours must match the data
for jj=2:N
ljpart(jj-1)=popweight(jj)*phi(jj)*lam(jj-1)^(-sigmaC/sigmaL)*(phi(jj)/phi(1))^(1/sigmaL);
end
l1ss=e_real/(popweight(1)*phi(1)+sum(ljpart))
L2(lam,l1ss,phi(2:N))
k1ss=(css^(-sigmaC)*css-omega*l1ss^(sigmaL)*l1ss)/(1-betta)/css^(-sigmaC)/(1+(alpha*(e_real/kss)^(1-alpha)-depr)*(1-tauK)); %lifetime budget constraint of group 1
kj=[k1ss kj2N] %capital of each group
kss %aggregate capital
popweight*kj'+kg %verify

cistst = [css lam*css]; %consumption of groups in the long run
listst = [l1ss, L2(lam,l1ss,phi(2:N))']; %hours worked of groups in the long run
cistst_sq = cistst
listst_sq = listst
%utilities at the status quo
if (sigmaC==1)
    Vsq = 1/(1-betta)*(log(cistst) - omega*listst.^(1+sigmaL)/(1+sigmaL));
else
    Vsq = 1/(1-betta)*((cistst.^(1-sigmaC)-1)/(1-sigmaC) - omega*listst.^(1+sigmaL)/(1+sigmaL));
end
Vsq

save('benchN2.mat')

rss = (betta^(-1)-1)/(1-tauK)+depr; %r from Euler at the steady state
(rss/alpha)^(1/(1-alpha))*kss % expre ss real total effective labor from r=F_k
govexp/kss^alpha/e_real^(1-alpha)
kss/kss^alpha/e_real^(1-alpha)
wss = (1-alpha)*kss^(alpha).*e_real^(-alpha); %w=F_e

tauKrev=tauK*(rss-depr)*kss
tauLrev=tauL*e_real*wss
tauKrev/(tauKrev+tauLrev)

